#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：python_learning 
@File ：yolo_test.py
@IDE  ：PyCharm 
@Author ：李涵彬
@Date ：2025/1/7 上午9:22
"""

import os
import subprocess
import torch
from typing import List


class YOLOv5Tester:
	def __init__(self, test_images_dir: str, output_dir: str, classes: List[str], use_gpu: bool = True):
		"""
		初始化YOLOv5测试器。

		:param test_images_dir: 包含测试图像的目录路径。
		:param output_dir: 输出模型测试结果的目录路径。
		:param classes: 类别列表。
		:param use_gpu: 是否使用GPU，默认为True。
		"""
		self.test_images_dir = os.path.abspath(test_images_dir)
		self.output_dir = os.path.abspath(output_dir)
		self.classes = classes

		# 选择设备
		self.device = 'cuda' if (use_gpu and torch.cuda.is_available()) else 'cpu'

	def test_model(self):
		"""
		测试YOLOv5模型。
		"""
		weights_path = os.path.join(self.output_dir, 'runs', 'train', 'exp', 'weights', 'best.pt')
		if not os.path.exists(weights_path):
			print(f"错误：模型权重文件 {weights_path} 不存在，请检查训练过程。")
			return

		output_dir = os.path.join(self.output_dir, 'runs', 'test')
		img_size = 640  # 图像大小

		# 测试命令
		yolov5_dir = os.path.dirname(os.path.abspath(__file__))  # 获取当前脚本目录
		test_command = [
			"python", os.path.join(yolov5_dir, 'yolov5', 'detect.py'),
			"--weights", weights_path,
			"--source", self.test_images_dir,
			"--img", str(img_size),
			"--device", self.device,
			"--project", output_dir
		]
		print("测试命令:", " ".join(test_command))
		result = subprocess.run(test_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
		print("Standard Output:", result.stdout.decode())
		print("Standard Error:", result.stderr.decode())
		if result.returncode != 0:
			print("测试失败，返回码:", result.returncode)
		else:
			print("测试成功")
